import pandas as pd
import numpy as np
import time
from sklearn.linear_model import LogisticRegression, LogisticRegressionCV
from sklearn.svm import SVC
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split
from sklearn.metrics import mean_squared_error, accuracy_score


def data_split(data: pd.DataFrame, start='2016-09', end='2018-10'):
    data.set_index(['date'], inplace=True)
    rtn_test = data.loc[end:]
    data.drop(['label', 'label_return', 'forward_return', 'name', 'code', 'industry'], axis=1, inplace=True)
    df_train = data.loc[start: end]
    df_test = data.loc[end:]
    df_train.reset_index()
    df_test.reset_index()
    return df_train.drop('label_cate', axis=1), df_test.drop('label_cate', axis=1), df_train['label_cate'], df_test['label_cate'], rtn_test.reset_index()


df = pd.read_csv('data_cate.csv')
df.fillna(0, inplace=True)
# print(np.where(np.array(df) == np.nan)[0])
# df = df.dropna()
# df.to_csv('test.csv')
# train_null = pd.isnull(df)
# print(len(train_null == True))
# train_null = df[train_null == True]
# print(train_null)

# Y_data = df['label_cate']
# Y_data = np.array(Y_data)
# X_data = df.drop(['label', 'label_return', 'label_cate', 'date', 'code', 'name'], axis=1)
#
# x_train,x_test,y_train,y_test = train_test_split(X_data, Y_data, test_size=0.3, random_state=None)
x_train, x_test, y_train, y_test, df_test = data_split(df)

# print(np.where(np.isnan(x_train))[0])
lr = LogisticRegression().fit(x_train, y_train)
prob = lr.predict_proba(x_test)[:, 1]
# pred = lr.predict(x_test)
print('LR:', accuracy_score(y_test, lr.predict(x_test)))

svm = SVC(probability=True).fit(x_train, y_train)
print('SVM:', accuracy_score(y_test, svm.predict(x_test)))
# prob = svm.predict_proba(x_test)[:, 1]

rf = RandomForestClassifier().fit(x_train, y_train)
print('RF:', accuracy_score(y_test, rf.predict(x_test)))
# prob = rf.predict_proba(x_test)[:, 1]

df_test['prob'] = pd.Series(prob)
df_test.to_csv('cate_test.csv', encoding='utf_8_sig')

